{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "uq6yF5pd6Fjj"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "##### Copyright 2023 Google LLC. SPDX-License-Identifier: Apache-2.0"
      ],
      "metadata": {
        "id": "uq6yF5pd6Fjj"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Copyright 2023 Google LLC. SPDX-License-Identifier: Apache-2.0\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at\n",
        "\n",
        "https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License."
      ],
      "metadata": {
        "id": "2mgX2mca6IP0"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **LLMs as General Pattern Machines:** PCFG Benchmark\n",
        "\n",
        "We observe that pretrained large language models (LLMs) are capable of autoregressively completing complex token sequences -- from arbitrary ones procedurally generated by probabilistic context-free grammars (PCFG), to more rich spatial patterns found in the Abstract Reasoning Corpus (ARC), a general AI benchmark, prompted in the style of ASCII art. Surprisingly, pattern completion proficiency can be partially retained even when the sequences are expressed using tokens randomly sampled from the vocabulary. These results suggest that without any additional training, LLMs can serve as general sequence modelers, driven by in-context learning. In this work, we investigate how these zero-shot capabilities may be applied to problems in robotics -- from extrapolating sequences of numbers that represent states over time to complete simple motions, to least-to-most prompting of reward-conditioned trajectories that can discover and represent closed-loop policies (e.g., a stabilizing controller for CartPole). While difficult to deploy today for real systems due to latency, context size limitations, and compute costs, the approach of using LLMs to drive low-level control may provide an exciting glimpse into how the patterns among words could be transferred to actions.\n",
        "\n",
        "Our **PCFG benchmark** is a procedurally generated, adjustable-difficulty benchmark for measuring abstract in-context learning sequence transformation capabilities in LLMs, based on the PCFG from [Hupkes et al. 2020](https://arxiv.org/abs/1908.08351). Here are the primitive operations in the PCFG that can be applied on one or two sequences of (arbitrary) tokens.\n",
        "\n",
        "<img src=\"https://socraticmodels.github.io/images/pcfg-ops.png\" height=\"300px\">\n",
        "\n",
        "These transformations include a collection of lexical rules that may be composed. The complexity of the transformations increase by increasing the number of tokens that represent the input sequences, or by increasing the number of primitives chained together. For example:\n",
        "\n",
        "<img src=\"https://socraticmodels.github.io/images/pcfg-example.png\" height=\"200px\">\n",
        "\n",
        "This colab runs GPT-3 on the PCFG benchmark with consistent tokenization (described more in Sec. 4 of the main paper). Evaluating on PCFG with random token alphabets can provide a more unbiased evaluation of pattern reasoning capabilities between LLMs. Our experiments suggest that PCFG completion accuracy improves with model scale. Note that we use PCFG for out-of-the-box in-context evaluation only (rather than for training data).\n",
        "\n",
        "### **Quick Start:**\n",
        "\n",
        "**Step 1.** Register for an [OpenAI API key](https://openai.com/blog/openai-api/) to use GPT-3 (there's a free trial) and enter it below\n",
        "\n",
        "**Step 2.** Menu > Runtime > Run all"
      ],
      "metadata": {
        "id": "r221djHjqHRL"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "openai_api_key = \"your-api-key-here\""
      ],
      "metadata": {
        "id": "Z77tvOok5SjP"
      },
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Setup**\n",
        "\n",
        "**Note:** only needs a CPU (public) runtime."
      ],
      "metadata": {
        "id": "itP9AeNSmfUQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install transformers\n",
        "!pip install sentencepiece\n",
        "!pip install openai\n",
        "!pip install tiktoken\n",
        "\n",
        "import numpy as np\n",
        "import tiktoken  # Faster than GPT2Tokenizer from HuggingFace.\n",
        "import openai\n",
        "import time\n",
        "\n",
        "openai.api_key = openai_api_key"
      ],
      "metadata": {
        "id": "NE5zQH9iooys"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **API:** Large Language Models\n",
        "\n",
        "Define helper functions to call large language models and the tokenizer.\n",
        "\n",
        "**Note:** this can get expensive. You can reduce the number of calls by reducing `num_datasets` when running on the benchmark."
      ],
      "metadata": {
        "id": "Yd-sk9BG6E1Y"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "encoding = tiktoken.get_encoding(\"gpt2\")\n",
        "encode = lambda s: encoding.encode(s, allowed_special={\"<|endoftext|>\"})\n",
        "decode = lambda s: encoding.decode(s)"
      ],
      "metadata": {
        "id": "5xYpucRmZEPw"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = \"text-davinci-003\"\n",
        "\n",
        "def LM(prompt, max_tokens=256, stop=None, temperature=0):\n",
        "  responses = openai.Completion.create(engine=model, prompt=prompt, max_tokens=max_tokens, temperature=temperature, stop=stop)\n",
        "  out_text = [response['text'] for response in responses['choices']]\n",
        "  return out_text"
      ],
      "metadata": {
        "id": "rNQmvEVUkP4X"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **PCFG:** Generator\n",
        "\n",
        "Define the rules that generate arbitrary sequence-to-sequence transformations. Each dataset is constructed using a PCFG rule applied on random tokens. The goal of in-context learning would be (i) to learn the PCFG rule from the dataset of examples and (ii) apply it to the new $x_{query}$. Both steps are done in a single forward pass."
      ],
      "metadata": {
        "id": "ODBGqT9jVo9L"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def pcfg_copy(x):\n",
        "  return x\n",
        "\n",
        "def pcfg_reverse(x):\n",
        "  return x[::-1]\n",
        "\n",
        "def pcfg_shift(x):\n",
        "  return x[1:] + x[:1]\n",
        "\n",
        "def pcfg_swap(x):\n",
        "  return x if len(x) < 2 else x[-1:] + x[1:-1] + x[:1]\n",
        "\n",
        "def pcfg_repeat(x):\n",
        "  return x + x\n",
        "\n",
        "def pcfg_echo(x):\n",
        "  return x + x[-1:]\n",
        "\n",
        "def pcfg_append(x, y):\n",
        "  return x + y\n",
        "\n",
        "def pcfg_prepend(x, y):\n",
        "  return y + x\n",
        "\n",
        "def pcfg_remove_first(x, y):\n",
        "  return y\n",
        "\n",
        "def pcfg_remove_second(x, y):\n",
        "  return x\n",
        "\n",
        "pcfg_fns = [pcfg_copy, pcfg_reverse, pcfg_shift, pcfg_swap, pcfg_repeat, pcfg_echo, pcfg_append, pcfg_prepend, pcfg_remove_first, pcfg_remove_second]\n",
        "def is_unary(pcfg_fn):\n",
        "  return pcfg_fn in [pcfg_copy, pcfg_reverse, pcfg_shift, pcfg_swap, pcfg_repeat, pcfg_echo]"
      ],
      "metadata": {
        "id": "nB89s2nGS-2m"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Rule:\n",
        "\n",
        "  def __init__(self, num_fns):\n",
        "    self.num_fns = num_fns\n",
        "    self.num_vars = 0\n",
        "\n",
        "    # Recursively generate a rule using a specific number of functions.\n",
        "    def gen_rule():\n",
        "      if self.num_fns == 0:\n",
        "        var_id = self.num_vars\n",
        "        self.num_vars += 1\n",
        "        return var_id\n",
        "      else:\n",
        "        rand_fn = np.random.choice(pcfg_fns)\n",
        "        self.num_fns -= 1\n",
        "        if is_unary(rand_fn):\n",
        "          return [rand_fn, gen_rule()]\n",
        "        else:\n",
        "          return [rand_fn, gen_rule(), gen_rule()]\n",
        "    self.fn_list = gen_rule()\n",
        "\n",
        "  # Apply the rule to a list of variables (lists).\n",
        "  def __call__(self, vars):\n",
        "    def apply_fns(params):\n",
        "      if type(params) == int:\n",
        "        return vars[params]\n",
        "      else:\n",
        "        next_fn = params[0]\n",
        "        if is_unary(next_fn):\n",
        "          return next_fn(apply_fns(params[1]))\n",
        "        else:\n",
        "          return next_fn(apply_fns(params[1]), apply_fns(params[2]))\n",
        "    return apply_fns(self.fn_list)"
      ],
      "metadata": {
        "id": "wiBlrpM8TDJ2"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **PCFG:** Datasets\n",
        "\n",
        "Procedurally generate datasets for infinite benchmarking."
      ],
      "metadata": {
        "id": "zGDTxmmXVvGB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "num_fns = 4  # Number of functions per PCFG rule. This increases task complexity.\n",
        "x_size = 10 # Number of tokens in each x.\n",
        "prompt_size = 1024  # Number of tokens in prompt.\n",
        "\n",
        "var_delim = \",\"  # Separates variables in each x.\n",
        "sample_delim = \"\\n---\\n\"  # Separates x from one another.\n",
        "input_output_delim = \"\\n\"  # Separates x and y from one another.\n",
        "\n",
        "num_datasets = 10  # How many datasets to generate?"
      ],
      "metadata": {
        "id": "i3ffFMj1TF-a"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def vars_to_tokens(vars):\n",
        "  tokens = []\n",
        "  for i, var in enumerate(vars):\n",
        "    token = encode(\" \" + str(var))\n",
        "    if i < len(vars) - 1:\n",
        "      tokens += token + encode(var_delim)\n",
        "    else:\n",
        "      tokens += token\n",
        "  return tokens\n",
        "\n",
        "def flatten(vars):\n",
        "  all_vars = []\n",
        "  for var in vars:\n",
        "    all_vars += var\n",
        "  return all_vars"
      ],
      "metadata": {
        "id": "b5B8nAUNGpIV"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Fixed test token vocabulary.\n",
        "vocab = [str(i) for i in range(10)]"
      ],
      "metadata": {
        "id": "l2LW_XvzWSa0"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Build PCFG datasets.\n",
        "pcfg_data = []\n",
        "for seed in range(num_datasets):\n",
        "  np.random.seed(seed)  # Reproduce generated datasets.\n",
        "\n",
        "  # Generate a random PCFG rule.\n",
        "  pcfg_rule = Rule(num_fns)\n",
        "\n",
        "  # Randomly split the list of tokens N-1 times to generate N variables for each x.\n",
        "  num_splits = pcfg_rule.num_vars - 1\n",
        "  split_idx = np.sort(np.random.choice(np.arange(1, x_size), size=num_splits, replace=False)).tolist()\n",
        "  split_idx = [0,] + split_idx + [x_size,]\n",
        "\n",
        "  # Fill the prompt with random examples generated using the PCFG rule.\n",
        "  prompt = []\n",
        "  context_size = 0\n",
        "  while True:\n",
        "\n",
        "    # Randomly generate a list of tokens (from vocab) to represent variables.\n",
        "    x_tokens = np.random.choice(vocab, size=x_size, replace=True).tolist()\n",
        "    x_vars = []\n",
        "    for i in range(1, len(split_idx)):\n",
        "      start, end = split_idx[i - 1], split_idx[i]\n",
        "      x_vars.append(x_tokens[start:end])\n",
        "\n",
        "    # Run the rule on x variables to get y.\n",
        "    y_vars = pcfg_rule(x_vars)\n",
        "    x_tokens = vars_to_tokens(flatten(x_vars)) + encode(input_output_delim)\n",
        "    y_tokens = vars_to_tokens(y_vars) + encode(sample_delim)\n",
        "\n",
        "    # Ensure the expected number of tokens does not exceed prompt size.\n",
        "    if len(prompt + x_tokens + y_tokens) > prompt_size:\n",
        "      break\n",
        "\n",
        "    prompt += x_tokens + y_tokens\n",
        "    context_size += 1\n",
        "    y_len = len(y_tokens)\n",
        "\n",
        "  prompt, target = prompt[:-y_len], prompt[-y_len:]\n",
        "  pcfg_data.append((prompt, target))\n",
        "\n",
        "print(\"Prompt:\\n\", decode(prompt))\n",
        "print(\"Target:\\n\", decode(target))\n",
        "print(\"Total number of in-context examples:\", context_size)"
      ],
      "metadata": {
        "id": "oAnR3Oo3TIUI",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "9826a58f-fa04-41bf-bcbc-49526097fccb"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Prompt:\n",
            "  5, 1, 0, 8, 8, 8, 2, 6, 8, 1\n",
            " 8, 1, 1\n",
            "---\n",
            " 8, 3, 5, 3, 6, 7, 9, 0, 8, 1\n",
            " 8, 1, 1\n",
            "---\n",
            " 8, 1, 6, 6, 2, 8, 4, 5, 3, 4\n",
            " 3, 4, 4\n",
            "---\n",
            " 0, 8, 0, 4, 5, 4, 8, 3, 8, 4\n",
            " 8, 4, 4\n",
            "---\n",
            " 8, 0, 1, 2, 3, 7, 7, 2, 0, 4\n",
            " 0, 4, 4\n",
            "---\n",
            " 6, 2, 2, 2, 3, 9, 2, 1, 3, 0\n",
            " 3, 0, 0\n",
            "---\n",
            " 1, 0, 2, 0, 4, 8, 0, 7, 0, 2\n",
            " 0, 2, 2\n",
            "---\n",
            " 4, 0, 1, 6, 7, 6, 1, 7, 6, 2\n",
            " 6, 2, 2\n",
            "---\n",
            " 5, 0, 0, 9, 5, 3, 7, 0, 6, 9\n",
            " 6, 9, 9\n",
            "---\n",
            " 2, 6, 5, 0, 9, 0, 5, 2, 5, 3\n",
            " 5, 3, 3\n",
            "---\n",
            " 6, 3, 1, 6, 1, 8, 8, 6, 8, 1\n",
            " 8, 1, 1\n",
            "---\n",
            " 7, 0, 1, 3, 6, 2, 0, 7, 4, 2\n",
            " 4, 2, 2\n",
            "---\n",
            " 8, 8, 1, 2, 4, 8, 1, 6, 4, 9\n",
            " 4, 9, 9\n",
            "---\n",
            " 3, 5, 1, 9, 2, 1, 7, 1, 8, 9\n",
            " 8, 9, 9\n",
            "---\n",
            " 7, 3, 4, 2, 3, 3, 1, 2, 6, 3\n",
            " 6, 3, 3\n",
            "---\n",
            " 1, 2, 4, 9, 6, 5, 7, 5, 1, 3\n",
            " 1, 3, 3\n",
            "---\n",
            " 6, 9, 4, 5, 2, 9, 5, 2, 9, 4\n",
            " 9, 4, 4\n",
            "---\n",
            " 8, 2, 8, 1, 0, 7, 9, 2, 9, 2\n",
            " 9, 2, 2\n",
            "---\n",
            " 7, 4, 2, 7, 3, 9, 1, 0, 0, 9\n",
            " 0, 9, 9\n",
            "---\n",
            " 4, 0, 3, 3, 6, 4, 8, 9, 0, 6\n",
            " 0, 6, 6\n",
            "---\n",
            " 9, 4, 7, 4, 0, 3, 8, 6, 5, 7\n",
            " 5, 7, 7\n",
            "---\n",
            " 3, 6, 2, 6, 5, 5, 6, 6, 5, 4\n",
            " 5, 4, 4\n",
            "---\n",
            " 7, 1, 8, 8, 6, 7, 4, 5, 3, 6\n",
            " 3, 6, 6\n",
            "---\n",
            " 0, 3, 2, 0, 6, 7, 8, 9, 3, 2\n",
            " 3, 2, 2\n",
            "---\n",
            " 1, 1, 9, 6, 4, 8, 5, 3, 3, 7\n",
            " 3, 7, 7\n",
            "---\n",
            " 3, 1, 1, 0, 1, 2, 0, 9, 0, 8\n",
            " 0, 8, 8\n",
            "---\n",
            " 6, 6, 5, 6, 8, 9, 9, 7, 3, 4\n",
            " 3, 4, 4\n",
            "---\n",
            " 4, 2, 7, 0, 8, 5, 8, 4, 2, 1\n",
            " 2, 1, 1\n",
            "---\n",
            " 7, 4, 0, 7, 5, 3, 3, 8, 0, 1\n",
            " 0, 1, 1\n",
            "---\n",
            " 2, 2, 6, 4, 6, 6, 9, 8, 4, 3\n",
            " 4, 3, 3\n",
            "---\n",
            " 4, 3, 6, 6, 1, 9, 7, 3, 5, 5\n",
            " 5, 5, 5\n",
            "---\n",
            " 5, 8, 2, 6, 4, 9, 7, 2, 2, 2\n",
            " 2, 2, 2\n",
            "---\n",
            " 6, 4, 4, 4, 1, 1, 1, 1, 5, 7\n",
            " 5, 7, 7\n",
            "---\n",
            " 0, 4, 5, 4, 7, 4, 7, 1, 9, 4\n",
            " 9, 4, 4\n",
            "---\n",
            " 1, 5, 3, 0, 3, 9, 9, 2, 8, 7\n",
            " 8, 7, 7\n",
            "---\n",
            " 4, 3, 1, 5, 3, 7, 9, 3, 1, 5\n",
            "\n",
            "Target:\n",
            "  1, 5, 5\n",
            "---\n",
            "\n",
            "Total number of in-context examples: 36\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Evaluate:** PCFG\n",
        "\n",
        "Run with GPT-3.\n"
      ],
      "metadata": {
        "id": "-WublGJUBn_1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print('x_size = ', x_size, 'num_fns = ', num_fns)\n",
        "verbose = True\n",
        "\n",
        "success = []\n",
        "count = 0\n",
        "for i, (prompt, target) in enumerate(pcfg_data):\n",
        "  count += 1\n",
        "\n",
        "  target_string = decode(target)\n",
        "  try:\n",
        "    output = LM(prompt, max_tokens=len(target_string)+3, stop=sample_delim)[0]\n",
        "  except:\n",
        "    print('*** failed call ***')\n",
        "    print('Prompt:', prompt)\n",
        "    print('***')\n",
        "    continue\n",
        "\n",
        "  success.append(target_string[:target_string.find(sample_delim)] in output)\n",
        "  print('Count:', count, ', Success percent:', int(100 * np.mean(success)))\n",
        "  if verbose and not success[-1]:\n",
        "    print('*** Failed: #', count)\n",
        "    print('Target:', target_string, 'Output:', output)\n",
        "    print('***')\n",
        "\n",
        "print('-'*10)\n",
        "print('Overall:')\n",
        "print('x_size = ', x_size, ', num_fns = ', num_fns, ', count =', count,\n",
        "      ', success rate =', np.mean(success))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4t-uUm96S6rQ",
        "outputId": "e2273f69-9184-493b-968a-de84140e37b5"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "x_size =  10 num_fns =  4\n",
            "Count: 1 , Success percent: 100\n",
            "Count: 2 , Success percent: 100\n",
            "Count: 3 , Success percent: 100\n",
            "Count: 4 , Success percent: 100\n",
            "Count: 5 , Success percent: 80\n",
            "*** Failed: # 5\n",
            "Target:  5, 9, 1, 1, 9, 3, 3\n",
            "---\n",
            " Output:  5, 9, 1, 1, 9, 5, 5\n",
            "***\n",
            "Count: 6 , Success percent: 83\n",
            "Count: 7 , Success percent: 71\n",
            "*** Failed: # 7\n",
            "Target:  8, 0, 5, 5, 8, 7, 0, 5, 5, 7\n",
            "---\n",
            " Output:  5, 0, 5, 5, 8, 7, 0, 5, 5, 7\n",
            "***\n",
            "Count: 8 , Success percent: 62\n",
            "*** Failed: # 8\n",
            "Target:  7, 1, 4, 9, 3, 6, 7, 1, 0, 7, 1, 4, 9, 3, 6, 7, 1, 0\n",
            "---\n",
            " Output:  1, 1, 4, 9, 3, 6, 7, 1, 7, 1, 1, 4, 9, 3, 6, 7, 1, 7\n",
            "***\n",
            "Count: 9 , Success percent: 55\n",
            "*** Failed: # 9\n",
            "Target:  1, 4, 1, 1, 7, 5, 4, 1, 1, 4, 1, 1, 7, 5, 4, 1\n",
            "---\n",
            " Output:  1, 1, 4, 1, 7, 5, 4, 1, 1, 1, 4, 1, 7, 5, 4, 1\n",
            "***\n",
            "Count: 10 , Success percent: 60\n",
            "----------\n",
            "Overall:\n",
            "x_size =  10 , num_fns =  4 , count = 10 , success rate = 0.6\n"
          ]
        }
      ]
    }
  ]
}